import {
  Flex,
  Box,
  Modal,
  ModalBody,
  ModalCloseButton,
  ModalContent,
  ModalHeader,
  ModalOverlay,
  Spacer,
  useDisclosure,
  SimpleGrid,
  useColorModeValue,
} from '@chakra-ui/react';
import { SelectItem } from '@mantine/core';
import IAIButton from 'common/components/IAIButton';
import { createSelector } from '@reduxjs/toolkit';
import { stateSelector } from 'app/store/store';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
import { MODEL_TYPE_MAP } from 'features/parameters/types/constants';
import { useCallback, useMemo } from 'react';
import { useGetControlNetModelsQuery } from 'services/api/endpoints/models';
// import { useTranslation } from 'react-i18next';
import { FaAngleRight } from 'react-icons/fa';
import {
  ControlNetConfig,
  controlNetModelChanged,
} from 'features/controlNet/store/controlNetSlice';
import { modelIdToControlNetModelParam } from 'features/parameters/util/modelIdToControlNetModelParam';
import { selectIsBusy } from 'features/system/store/systemSelectors';
import { forEach, map } from 'lodash-es';
import EasyControlNetItem, { CONTROLNETMODES } from './EasyControlNetItem';
type ParamControlNetModelProps = {
  controlNet: ControlNetConfig;
};
const selector = createSelector(
  stateSelector,
  ({ generation }) => {
    const { model } = generation;
    return { mainModel: model };
  },
  defaultSelectorOptions
);
const EasyControlNetSelect = (props: ParamControlNetModelProps) => {
  const dispatch = useAppDispatch();
  const { isOpen, onOpen, onClose } = useDisclosure();
  const btn_bg = useColorModeValue('base.800', 'base.750');
  const { controlNetId, model: controlNetModel, isEnabled } = props.controlNet;
  const isBusy = useAppSelector(selectIsBusy);

  const { mainModel } = useAppSelector(selector);

  const { data: controlNetModels } = useGetControlNetModelsQuery();

  const data = useMemo(() => {
    if (!controlNetModels) {
      return [];
    }

    const data: SelectItem[] = [];

    forEach(controlNetModels.entities, (model, id) => {
      if (!model) {
        return;
      }

      const disabled = model?.base_model !== mainModel?.base_model;

      data.push({
        value: id,
        label: model.model_name,
        group: MODEL_TYPE_MAP[model.base_model],
        disabled,
        tooltip: disabled
          ? `Incompatible base model: ${model.base_model}`
          : undefined,
        ...model,
      });
    });

    return data;
  }, [controlNetModels, mainModel?.base_model]);

  // grab the full model entity from the RTK Query cache
  const selectedModel = useMemo(
    () =>
      controlNetModels?.entities[
        `${controlNetModel?.base_model}/controlnet/${controlNetModel?.model_name}`
      ] ?? null,
    [
      controlNetModel?.base_model,
      controlNetModel?.model_name,
      controlNetModels?.entities,
    ]
  );

  const handleModelChanged = useCallback(
    (v: string | null) => {
      if (!v) {
        return;
      }

      const newControlNetModel = modelIdToControlNetModelParam(v);

      if (!newControlNetModel) {
        return;
      }

      dispatch(
        controlNetModelChanged({ controlNetId, model: newControlNetModel })
      );
      onClose();
    },
    [controlNetId, dispatch, onClose]
  );
  const modelMode = selectedModel
    ? CONTROLNETMODES[
        selectedModel.model_name as keyof typeof CONTROLNETMODES
      ] || selectedModel.model_name
    : undefined;
  return (
    <>
      <Flex alignItems="center" py={1} flexGrow={1} gap="1">
        <Box fontSize="sm" fontWeight="600">
          控制类型
        </Box>
        <Spacer />
        <IAIButton
          aria-label="选择控制类型"
          tooltip="选择控制类型"
          rightIcon={<FaAngleRight />}
          size="xs"
          rounded="base"
          onClick={onOpen}
          bg={btn_bg}
          _hover={{
            bg: 'base.750',
          }}
          isDisabled={isBusy || !isEnabled}
        >
          {selectedModel ? modelMode : '选择'}
        </IAIButton>
      </Flex>

      <Modal
        scrollBehavior="inside"
        isOpen={isOpen}
        onClose={onClose}
        size={{ base: 'xl', md: '3xl', lg: '4xl', xl: '5xl', '2xl': '5xl' }}
        isCentered
      >
        <ModalOverlay />
        <ModalContent bg="base.950" fontFamily="Inter">
          <ModalCloseButton />
          <ModalHeader fontWeight="bold" py={2}>
            选择控制类型
          </ModalHeader>
          <ModalBody>
            <Flex
              justify="center"
              padding="0 0 1.5rem 0"
              width="100%"
              height="100%"
            >
              <SimpleGrid columns={{ base: 2, md: 3 }} spacing={4}>
                {map(data, (c) => (
                  <EasyControlNetItem
                    key={c.label}
                    {...c}
                    isSelected={selectedModel?.model_name === c.label}
                    onClick={handleModelChanged}
                  />
                ))}
              </SimpleGrid>
            </Flex>
          </ModalBody>
        </ModalContent>
      </Modal>
    </>
  );
};

export default EasyControlNetSelect;
